# 基本思路：
# 使用https://www.luogu.com.cn/lg4/captcha看验证码
# 然后使用https://www.luogu.com.cn/do-auth/password和{"username": "","password": "","captcha" :  ""}登录
# 拿到cookie中的__client_id和_uid.然后想要测评哪道题就请求题目页面，拿到该题目页面的csrf-token（会写在HTML响应里）
# 然后Post该接口https://www.luogu.com.cn/fe/api/problem/submit/P1000,头部里携带刚刚获取到的x-csrf-token和referer以及cookie
# 请求体里装代码{"code": "#include <bits/stdc++.h>\nusing namespace std;int main(){cout<<\"Hello\"<<endl;}","lang": 0,"enableO2": 0}
# 接下来就能提交测评了，响应会是一个rid，然后将rid放进该接口进行请求（携带cookie）就能得到测评结果https://www.luogu.com.cn/record/208409537?_contentOnly=1
# 这块可以做一个只请求一次的版本或和一个轮询的版本（直到请求到测评完成的响应停止）

import requests
import sys
import time
import re
import json
from PIL import Image
from io import BytesIO

class LuoguLogin:
    def __init__(self):
        """初始化 Session 以维持登录状态"""
        self.session = requests.Session()
        self.base_url = "https://www.luogu.com.cn"
        self.headers = {
            # 避免被识别为爬虫，需要带一个模拟浏览器的 User-Agent
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/120.0.0.0 Safari/537.36"
            )
        }

    def get_captcha(self):
        """获取验证码并在终端显示，让用户手动输入"""
        captcha_url = f"{self.base_url}/lg4/captcha"
        response = self.session.get(captcha_url, headers=self.headers)

        if response.status_code == 200:
            # 使用 PIL 显示验证码
            img = Image.open(BytesIO(response.content))
            img.show()
            # 让用户手动输入验证码
            captcha_code = input("请输入验证码: ")
            return captcha_code
        else:
            print("获取验证码失败，状态码:", response.status_code)
            sys.exit(1)

    def login(self, username, password):
        """执行登录"""
        captcha_code = self.get_captcha()
        login_url = f"{self.base_url}/do-auth/password"
        payload = {
            "username": username,
            "password": password,
            "captcha": captcha_code
        }

        response = self.session.post(login_url, json=payload, headers=self.headers)
        if response.status_code == 200:
            result = response.json()
            if "errorCode" in result:
                print(f"登录失败: {result.get('errorMessage')}")
                return False
            else:
                print(f"登录成功，用户名: {result.get('username')}")
                # 可选：保存Cookie
                self.save_session_cookies()
                return True
        else:
            print("登录请求失败，状态码:", response.status_code)
            return False

    def save_session_cookies(self):
        """将当前 Session 中的 Cookie 保存到本地文件"""
        with open("cookies.txt", "w", encoding="utf-8") as f:
            for cookie in self.session.cookies:
                f.write(f"{cookie.name}={cookie.value}\n")

    def load_session_cookies(self):
        """从本地文件加载 Cookie，写入 Session"""
        try:
            with open("cookies.txt", "r", encoding="utf-8") as f:
                cookies = f.readlines()
            for cookie in cookies:
                name, value = cookie.strip().split("=")
                self.session.cookies.set(name, value)
            print("已加载登录 Cookie")
        except FileNotFoundError:
            print("未找到已保存的 Cookie，需重新登录")

    def check_login_status(self):
        check_url = f"{self.base_url}/api/feed/watching"
        resp = self.session.get(check_url, headers=self.headers)
        if resp.status_code == 200:
            # 如果能够返回 200 并且响应体里含有 data 列表等字段，就说明已登录
            print("当前处于登录状态")
            return True
        else:
            print("登录状态已失效或无法确认，返回码:", resp.status_code)
            return False

    def get_csrf_token_from_problem_page(self, pid):
        """访问题目页面，解析 HTML 中的 CSRF Token"""
        page_url = f"{self.base_url}/problem/{pid}"
        resp = self.session.get(page_url, headers=self.headers)
        if resp.status_code != 200:
            print(f"获取题目页面失败，状态码: {resp.status_code}")
            return None

        html_text = resp.text
        # 1) 尝试通过 window._csrf = "xxxx"
        match = re.search(r'window._csrf\s*=\s*"([^"]+)"', html_text)
        if match:
            return match.group(1)
        # 2) 或者尝试 <meta name="csrf-token" content="xxxx">
        match = re.search(r'<meta name="csrf-token"\s+content="([^"]+)"', html_text)
        if match:
            return match.group(1)

        print("题目页面中未找到 CSRF Token，可能该题不需或已改版")
        return None

    def submit_code(self, pid, code, language=0, enable_o2=False):
        """
        提交代码到洛谷
        :param pid: 题号 (e.g. P1000)
        :param code: 源代码字符串
        :param language: 语言 ID (0 可能对应 C++，具体需查对应语言列表)
        :param enable_o2: 是否启用 O2 优化
        :return: 提交记录 ID (rid) 或 None
        """
        csrf_token = self.get_csrf_token_from_problem_page(pid)
        if not csrf_token:
            print("无法获取或不需要CSRF Token，但若提交失败需再检查")
            # 有时可能不需要 Token，如果报错还是 Invalid CSRF 就要再来查

        submit_url = f"{self.base_url}/fe/api/problem/submit/{pid}"
        payload = {
            "code": code,
            "lang": language,
            "enableO2": int(enable_o2)
        }

        # 补充提交时所需的头部
        headers = self.headers.copy()
        if csrf_token:
            headers["x-csrf-token"] = csrf_token
        headers["Referer"] = f"{self.base_url}/problem/{pid}"

        # 提交代码
        resp = self.session.post(submit_url, json=payload, headers=headers)
        if resp.status_code == 200:
            try:
                data = resp.json()
                if "rid" in data:
                    print(f"提交成功，rid = {data['rid']}")
                    return data["rid"]
                else:
                    print("提交失败，后端返回:", data)
            except Exception as e:
                print("解析提交结果 JSON 失败:", e)
        else:
            print(f"提交失败, HTTP状态码: {resp.status_code}, 响应: {resp.text}")
        return None

    def get_judging_result(self, rid, max_poll=10):
        """
        轮询获取评测结果
        :param rid: 提交记录 ID
        :param max_poll: 轮询次数上限
        :return: 评测结果的 JSON 或 None
        """
        # 加上 _contentOnly=1 获取 JSON 格式
        result_url = f"{self.base_url}/record/{rid}?_contentOnly=1"

        for i in range(max_poll):
            resp = self.session.get(result_url, headers=self.headers)
            if resp.status_code == 200:
                try:
                    data = resp.json()
                    record = data["currentData"]["record"]
                    # record["status"] 状态值不一定是 0=成功，可以是 14=WA 等
                    # 只要不是还在Judging就算有结果
                    status_code = record["status"]
                    if status_code in (0, 12, 14, 15, 16, 17, 18):
                        # 0=评测中 or 12=AC or 14=WA ...
                        # 不同数字对应不同评测结果
                        print(f"轮询第{i+1}次，status = {status_code}")
                        # 如果需要进一步判断是否“完成”可看 subtask 测试完成情况
                    else:
                        print(f"轮询第{i+1}次，未知的 status = {status_code}")
                    return data
                except Exception as e:
                    print("解析评测结果 JSON 失败:", e)
                    return None
            else:
                print(f"HTTP状态码: {resp.status_code}, 无法获取评测结果")
                return None
            # 如果还想多次轮询就等待 3 秒
            time.sleep(3)

        print("轮询次数已达上限，可能评测排队过长")
        return None


if __name__ == "__main__":
    luogu = LuoguLogin()

    # 尝试载入Cookie（若之前已登录过）
    luogu.load_session_cookies()

    # 如果没有登录或失效则重新登录
    if not luogu.check_login_status():
        username = input("请输入洛谷用户名: ")
        password = input("请输入洛谷密码: ")
        if not luogu.login(username, password):
            print("登录失败，程序退出")
            sys.exit()

    print("已确认处于登录状态，开始提交示例题目")

    # 准备提交示例
    pid = "P1000"
    sample_code = r"""#include <bits/stdc++.h>
    using namespace std;
    int main(){
        cout<<"Hello World"<<endl;
        return 0;
    }
    """

    # 提交
    rid = luogu.submit_code(pid, sample_code, language=0, enable_o2=False)
    if rid:
        print("开始查询评测结果...")
        res = luogu.get_judging_result(rid, max_poll=10)
        print("最终评测结果：", res)
    else:
        print("提交失败，无法查询评测结果。")
